import copy
import time
import torch
import torch.nn as nn
from model.full_model import Translayout
from boxex_utils.data_utils import data_loader
from utils.parse_config import get_config
from torch.utils.data import DataLoader


def train_epoch(model, optimizer,config):

    model.train()
    losses = 0
    train_dataloader = data_loader(config.file_dir,config.batch_size,config.max_len)
    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    loss_fn = torch.nn.CrossEntropyLoss(ignore_index=0)


    for boxes_array,token_array,positions_array,labels_array in train_dataloader:
        boxes_array = boxes_array.to(DEVICE)
        token_array = token_array.to(DEVICE)
        positions_array = positions_array.to(DEVICE)
        labels_array = labels_array.to(DEVICE)
        # boxes_input = boxes_array[:-1, :]
        # tokens_input = token_array[:-1,:]
        boxes_input = boxes_array
        tokens_input = token_array
        position_input = positions_array
        labels = labels_array

        logits = model(tokens_input, boxes_input, position_input)
        _, preds = torch.max(logits, 1)
        optimizer.zero_grad()
        loss = loss_fn(logits.reshape(-1, logits.shape[-1]), labels.reshape(-1))
        loss.backward()

        optimizer.step()
        losses += loss.item()

    return losses / len(train_dataloader)


def evaluate(model,config):
    model.eval()
    losses = 0
    running_corrects=0
    val_dataloader = data_loader(config.file_dir,config.batch_size,config.max_len)
    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    for boxes_array,token_array,positions_array,labels_array in val_dataloader:
        boxes_array = boxes_array.to(DEVICE)
        token_array = token_array.to(DEVICE)
        position_array = positions_array.to(DEVICE)
        labels_array = labels_array.to(DEVICE)
        # boxes_input = boxes_array[:-1, :]
        # tokens_input = token_array[:-1,:]
        boxes_input = boxes_array
        tokens_input = token_array
        position_input = position_array
        labels = labels_array
        logits = model(tokens_input, boxes_input, position_input)
        _, preds = torch.max(logits, 1)

        loss = loss_fn(logits.reshape(-1, logits.shape[-1]), labels.reshape(-1))

        losses += loss.item()
        running_corrects += torch.sum(preds == labels)
    return losses / len(val_dataloader)

if __name__=="__main__":
    from timeit import default_timer as timer
    NUM_EPOCHS = 18
    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    ubuntu_file = "/mnt/myproject/classfication/translayout/utils/config.json"
    os_file = "/Users/qianghaozhang/Documents/PythonProject/JXLZ/classification/translayout/utils/config.json"
    config = get_config(os_file)
    model = Translayout(config).to(DEVICE)
    loss_fn = torch.nn.CrossEntropyLoss(ignore_index=0)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9)
    for epoch in range(1, NUM_EPOCHS+1):
        start_time = timer()
        train_loss = train_epoch(model, optimizer,config)
        end_time = timer()
        val_loss = evaluate(model,config)
        print((f"Epoch: {epoch}, Train loss: {train_loss:.3f}, Val loss: {val_loss:.3f}, "f"Epoch time = {(end_time - start_time):.3f}s"))